import os
import time
import asyncio
import json
import logging
import subprocess
import uuid

from aiohttp import web, ClientSession
from logging.handlers import RotatingFileHandler

log_dir = "logs"
os.makedirs(log_dir, exist_ok=True)

# 创建日志记录器
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')

console_handler = logging.StreamHandler()
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)

file_handler = RotatingFileHandler(
    os.path.join(log_dir, "proxy—main.log")
)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)

# 配置参数
PROXY_IP = "0.0.0.0"
PROXY_PORT = 20011


class MCPProxyService:
    def __init__(self, mcp_command, sse_server_url, sse_client_id, env_vars=None, timeout=10):
        self.mcp_process = None
        self.mcp_command = mcp_command
        self.sse_server_url = sse_server_url
        self.sse_client_id = sse_client_id
        self.env_vars = env_vars or {}  # 环境变量字典
        self.sse_session = None
        self.app = web.Application()
        self.setup_routes()
        self.response_queue = asyncio.Queue()  # 新增响应队列
        self.read_lock = asyncio.Lock()  # 保留锁，但仅用于命令执行
        self.timeout = timeout

    def setup_routes(self):
        self.app.router.add_get('/health', self.health_check)
        self.app.router.add_post('/execute', self.execute_mcp_command)

    async def execute_mcp_command(self, request):
        """执行MCP命令并返回结果"""
        try:
            # 从请求中获取MCP命令和参数
            data = await request.json()
            command = data.get('command')  # 例如: "tools/list"
            params = data.get('params', {})  # 命令参数

            if not command:
                return web.json_response({"error": "Missing 'command' parameter"}, status=400)

            # 构造完整的MCP命令（根据MCP协议调整格式）
            full_command, request_id = self.format_mcp_command(command, params)

            # 发送命令到MCP Server
            if self.mcp_process.stdin:
                self.mcp_process.stdin.write(f"{full_command}\n".encode('utf-8'))
                await self.mcp_process.stdin.drain()

                # 等待并获取MCP响应
                async with self.read_lock:  # 使用锁确保同一时间只有一个协程可以读取数据
                    response = await self.wait_for_mcp_response(command, request_id)
                return web.json_response({"result": response})
            else:
                return web.json_response({"error": "MCP进程未运行"}, status=500)

        except Exception as e:
            logger.error(f"执行MCP命令失败: {str(e)}")
            return web.json_response({"error": str(e)}, status=500)

    def format_mcp_command(self, command, params):
        """格式化为JSON-RPC请求"""
        request_id = str(uuid.uuid4())  # 生成唯一ID
        return json.dumps({
            "jsonrpc": "2.0",
            "method": command,
            "params": params,
            "id": request_id
        }), request_id

    async def wait_for_mcp_response(self, command, request_id):
        """从响应队列获取并解析MCP Server的JSON响应"""
        try:
            timeout = self.timeout  # 超时时间
            start_time = time.time()

            while time.time() - start_time < timeout:
                try:
                    # 从队列获取响应（带超时）
                    message = await asyncio.wait_for(self.response_queue.get(), timeout=0.5)
                    logger.info(message)

                    try:
                        response = json.loads(message)
                        if response.get('id') == request_id:  # 用id匹配
                            return response
                        # 不匹配的响应放回队列或处理
                        await self.response_queue.put(message)
                    except json.JSONDecodeError:
                        continue

                except asyncio.TimeoutError:
                    continue  # 继续等待，直到整体超时

            raise Exception(f"MCP响应超时 ({timeout}秒)")

        except Exception as e:
            logger.error(f"获取MCP响应失败: {str(e)}")
            raise

    def is_response_complete(self, responses):
        """判断响应是否已完整（需要根据MCP Server的响应格式调整）"""
        try:
            # 尝试将所有响应合并为一个JSON字符串并解析
            response_str = ''.join(responses)
            response_json = json.loads(response_str)
            # 检查是否包含必要的字段
            if 'jsonrpc' in response_json and 'id' in response_json:
                return True
        except json.JSONDecodeError:
            pass
        return False

    async def monitor_mcp_process(self):
        """监控MCP进程，意外退出时重启"""
        while True:
            if self.mcp_process and self.mcp_process.returncode is not None:
                logger.error(f"MCP进程意外退出，返回码: {self.mcp_process.returncode}，尝试重启...")
                await self.start_mcp_process()
            await asyncio.sleep(5)  # 检查间隔

    async def start_mcp_process(self):
        """启动MCP Server进程并设置stdio管道"""
        # 合并默认环境变量和自定义环境变量
        full_env = {
            "HOME": ".",
            "LOGNAME": "ifeng",
            "PATH": "/Users/ifeng/.npm/_npx/5a9d879542beca3a/node_modules/.bin:/Users/ifeng/node_modules/.bin:/Users/node_modules/.bin:/node_modules/.bin:/usr/local/lib/node_modules/npm/node_modules/@npmcli/run-script/lib/node-gyp-bin:/Library/Java/JavaVirtualMachines/jdk-17.jdk/Contents/Home/bin:/usr/local/mysql/bin:/Library/Frameworks/Python.framework/Versions/3.12/bin:/usr/local/bin:/System/Cryptexes/App/usr/bin:/usr/bin:/bin:/usr/sbin:/sbin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/local/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/bin:/var/run/com.apple.security.cryptexd/codex.system/bootstrap/usr/appleinternal/bin:/opt/homebrew/bin:/usr/local/go/bin",
            "SHELL": "/bin/zsh",
            "TERM": "xterm-256color",
            "USER": "ifeng",
            "UP_ACCESS_TYPE": '01',
            "UP_MER_ID": "777290058211279",
            "UP_TOKEN_TYPE": "01",
            "UP_SIGN_CERT_PATH": "./cert/acp_test_sign.pfx",
            "UP_SIGN_CERT_PWD": "000000",
            "UP_VALIDATE_CERT_DIR": "./cert/verify_sign_acp.cer"
        }
        self.mcp_process = await asyncio.create_subprocess_shell(
            "unionpay-mcp-server",
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            env=full_env  # 传递环境变量
        )
        logger.info(f"MCP Server 进程已启动，PID: {self.mcp_process.pid}")

        # 启动异步任务读取MCP输出
        asyncio.create_task(self.read_mcp_output())
        asyncio.create_task(self.read_mcp_errors())
        asyncio.create_task(self.monitor_mcp_process())

    async def read_mcp_output(self):
        """从MCP Server读取输出并放入响应队列"""
        output = []
        while True:
            if self.mcp_process.stdout is None:
                await asyncio.sleep(1)
                continue

            line = await self.mcp_process.stdout.readline()
            if not line:
                logger.warning("MCP进程输出流已关闭")
                break

            message = line.decode('utf-8').strip()
            output.append(message)
            logger.info(f"从MCP收到: {message}")

            # 将消息放入响应队列
            await self.response_queue.put(message)

            # 发送到SSE服务器
            # await self.send_to_sse_server(message)

        # 进程结束后，保存完整输出
        self.last_mcp_output = {
            'stdout': output,
            'returncode': self.mcp_process.returncode
        }
        logger.info(f"MCP进程已结束，返回码: {self.mcp_process.returncode}")

    async def read_mcp_errors(self):
        """读取MCP Server的错误输出"""
        errors = []  # 存储错误信息
        while True:
            if self.mcp_process.stderr is None:
                await asyncio.sleep(1)
                continue

            line = await self.mcp_process.stderr.readline()
            if not line:
                break

            error = line.decode('utf-8').strip()
            errors.append(error)  # 收集错误
            logger.error(f"MCP错误: {error}")

        # 保存错误信息
        if hasattr(self, 'last_mcp_output'):
            self.last_mcp_output['stderr'] = errors

    async def health_check(self, request):
        """健康检查接口"""
        status = "running" if self.mcp_process and self.mcp_process.returncode is None else "stopped"
        return web.json_response({"status": status})

    async def cleanup(self):
        """清理资源"""
        if self.mcp_process and self.mcp_process.returncode is None:
            self.mcp_process.terminate()
            await self.mcp_process.wait()

        if self.sse_session:
            await self.sse_session.close()


async def main():
    # 配置参数
    mcp_command = "unionpay-mcp-server"  # 替换为实际的MCP启动命令
    # sse_server_url = "http://127.0.0.1:8000/sse"  # 替换为实际的SSE服务端URL
    sse_server_url = "https://c75c42325b30.ngrok-free.app/sse"  # 替换为实际的SSE服务端URL
    sse_client_id = "mcp-proxy-client-001"  # 客户端唯一标识

    proxy = MCPProxyService(mcp_command, sse_server_url, sse_client_id)

    # 设置关闭处理
    async def on_shutdown(app):
        await proxy.cleanup()

    proxy.app.on_shutdown.append(on_shutdown)

    # 启动服务
    await proxy.start_mcp_process()

    # 启动Web服务器
    runner = web.AppRunner(proxy.app)
    await runner.setup()
    site = web.TCPSite(runner, PROXY_IP, PROXY_PORT)
    await site.start()

    logger.info("MCP代理服务已启动，监听端口20011")

    # 保持运行
    await asyncio.Event().wait()


if __name__ == "__main__":
    asyncio.run(main())
